package com.neihe.common.utils;

import com.neihe.common.constant.Const;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

/**
 * 限流通用工具
 * 针对于不同的情况进行限流操作，支持限流升级
 */
@Slf4j
@Component
public class FlowUtils {

    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 针对于单次频率限制，请求成功后，在冷却时间内不得再次进行请求，如3秒内不能再次发起请求
     *
     * @param key       键
     * @param blockTime 限制时间
     * @return 是否通过限流检查
     */
    public boolean limitOnceCheck(String key, int blockTime) {
        return internalCheck(key, 1, blockTime, (overclock) -> false);
    }

    /**
     * 针对IP地址进行邮件验证码获取限流
     *
     * @param address 地址
     * @return true：通过；false：未通过
     */
    public boolean verifyLimit(String address) {
        String key = Const.LIMIT_EMAIL_CODE + address;
        return limitOnceCheck(key, Const.LIMIT_EMAIL_CODE_TIME);
    }

    /**
     * 内部使用请求限制主要逻辑
     *
     * @param key       计数键
     * @param frequency 请求频率
     * @param period    计数周期
     * @param action    限制行为与策略
     * @return 是否通过限流检查
     */
    private boolean internalCheck(String key, int frequency, int period, LimitAction action) {
        String count = redisTemplate.opsForValue().get(key);
        if (count != null) {
            long value = Optional.ofNullable(redisTemplate.opsForValue().increment(key)).orElse(0L);
            int c = Integer.parseInt(count);
            if (value != c + 1) redisTemplate.expire(key, period, TimeUnit.SECONDS);
            return action.run(value > frequency);
        } else {
            redisTemplate.opsForValue().set(key, "1", period, TimeUnit.SECONDS);
            return true;
        }
    }

    /**
     * 内部使用，限制行为与策略
     */
    private interface LimitAction {
        boolean run(boolean overclock);
    }
}
